/* Copyright (C) 2005, 2006, 2007, 2008, 2009, 2010 Mario Orsi
   This file is part of Brahms.
   Brahms is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
   Brahms is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
   You should have received a copy of the GNU General Public License along with Brahms.  If not, see <http://www.gnu.org/licenses/>. */

// latDiff.c -- Direct evaluation of lipid lateral diffusion, adapted from [Rapaport p. 124].

#include "dataStructs.h"

/* external data references */
extern VecRTwo *xyLipidCom;
extern const VecR *siteCoords;
extern VecRTwo *xyLipidCom, *xyLipidComTrue;
extern VecRTwo xySoluteCom, xySoluteComTrue;
extern const Site *site;  
extern const VecR region; 
extern const int nLipids, stepCount;
extern const real deltaT, timeNow;
extern const LipidStruct dopcStruct, dopeStruct;
extern const int nSites;

LatDiffBuff *latDiffBuff; // buffer for time-dependent measurement 
real *latSqrDiffAv; // diffusion mean 
int limitLatDiffAv; // total number of individual estimates used to produce the average value of the lateral diff coeff   
int nBuffLatDiff; // number of sets of data being collected at any time, each occupying a separate storage buffer 
int nValLatDiff; // number of measurements contributing to the set used to produce a single unaveraged estimate of the lat diff coeff
int stepLatDiff; // measurement frequency (number of steps between measurements)
static int countLatDiffAv; // number of measurements in set 

static void PrintLatDiff()
{
  real tVal;
  int j;

  FILE *fp;
  char latDiffFile[ 20 ];
  char time_ns[ 9 ]; // string containing timeNow, e.g. "100", in ns

  //  sprintf( time_ns, "%d", stepCount ); // 
  sprintf( time_ns, "%d", Nint( timeNow *  TIME_ns ) );
  latDiffFile[ 0 ] = '\0';
  strcat( latDiffFile, "latDiff-" );
  strcat( latDiffFile, time_ns );
  strcat( latDiffFile, ".dat" );
  fp = fopen( latDiffFile, "w" );

  for ( j = 0; j < nValLatDiff; j++ ) {
    tVal = j * stepLatDiff * deltaT;
    //fprintf( fp, "%17.4f %15.10f\n", tVal * TIME_ps, latSqrDiffAv[ j ] * Sqr( LENGTH_CM ) / TIME_s * 1e5); /* [ps], [10^(-5)cm^2/s] */
    //fprintf( fp, "%17.4f %15.10f\n", tVal * TIME_ps, latSqrDiffAv[ j ] * Sqr( LENGTH_CM ) / TIME_s * 1e8); /* [ps], [10^(-8)cm^2/s] */
    //    fprintf( fp, "%12.5f %12.3f\n", tVal * TIME_ns, latSqrDiffAv[ j ] * Sqr( LENGTH_CM ) / TIME_s * 1e8); /* [ns], [10^(-8)cm^2/s] */
    fprintf( fp, "%12.5f %12.3f\n", tVal * TIME_ns, latSqrDiffAv[ j ] * Sqr( LENGTH_nm ) / TIME_mus ); /* [ns], [nm^2/microsecond] */
    //  fflush( fp );
  }
  fclose( fp );
}

static void ZeroLatDiff()
{
  int j;

  countLatDiffAv = 0;
  for ( j = 0; j < nValLatDiff; j++ ) latSqrDiffAv[ j ] = 0.;  /* resetting the calculation */
}


/* InitLatDiff. Initialises the calculation. Called from SetUpJob(). 
   The values initially assigned to "count" determine the spacing between measurements: negative initial values delay
   the start of data collection for each set of measurements until the appropriate moment. */
void InitLatDiff()
{
  int nb;

  for ( nb = 0; nb < nBuffLatDiff; nb++ ) {
    latDiffBuff[ nb ].count = - nb * nValLatDiff / nBuffLatDiff;
  }
  ZeroLatDiff();
}

static void AccumLatDiff()
{
  real fac;
  int j, nb;
  
  for ( nb = 0; nb < nBuffLatDiff; nb++ ) {
    if ( latDiffBuff[ nb ].count == nValLatDiff ) { /* if buffer "nb" is full */
      for ( j = 0; j < nValLatDiff; j++ ) { /* scan through the buffer elements */
	latSqrDiffAv[ j ] += latDiffBuff[ nb ].latSqrDisp[ j ]; /* accumulate LatDiff */
      }
      latDiffBuff[ nb ].count = 0;  /* reset buffer counter */
      ++countLatDiffAv;
      if ( countLatDiffAv == limitLatDiffAv ) { /* if the #individual estimates required has been reached */
        fac = 1. / ( 4. * nLipids * stepLatDiff * deltaT * limitLatDiffAv ); // 4 <- (2 * D), being D = 2 in this case (lipid diffusion in the xy plane) 
        for ( j = 1; j < nValLatDiff; j++ ) {
	  latSqrDiffAv[ j ] *= fac / j;
	}
	PrintLatDiff();
	ZeroLatDiff();
      }
    }
  }
}

/* Called by SingleStep() every "stepLatDiff" steps (after equilibration, obviously). */
void EvalLatDiff()
{
  VecRTwo latDisp; // lateral lipid displacement (current - initial) in the xy plane 
  int n, nb, ni;

  for ( nb = 0; nb < nBuffLatDiff; nb++ ) { /* scan through the individual buffers */
    if ( latDiffBuff[ nb ].count == 0 ) { /* if empty buffer */
      DO_LIPID { // set origin vectors
        latDiffBuff[ nb ].orgTwo[ n ] = xyLipidCom[ n ];
        latDiffBuff[ nb ].latDispTrueTwo[ n ] = xyLipidCom[ n ];
      }
    }
    if ( latDiffBuff[ nb ].count >= 0 ) {
      ni = latDiffBuff[ nb ].count;
      latDiffBuff[ nb ].latSqrDisp[ ni ] = 0.;
      DO_LIPID {
        VSubTwo( latDisp, latDiffBuff[ nb ].latDispTrueTwo[ n ], xyLipidCom[ n ] );
        VDivTwo( latDisp, latDisp, region );
        latDisp.x = Nint( latDisp.x );
        latDisp.y = Nint( latDisp.y );
        VMulTwo( latDisp, latDisp, region );
        VAddTwo( latDiffBuff[ nb ].latDispTrueTwo[ n ], xyLipidCom[ n ], latDisp );
        VSubTwo( latDisp, latDiffBuff[ nb ].latDispTrueTwo[ n ], latDiffBuff[ nb ].orgTwo[ n ] );
        latDiffBuff[ nb ].latSqrDisp[ ni ] += VLenSqTwo( latDisp );
      }
    }
    ++latDiffBuff[ nb ].count;  /* increase buffer counter */
  }
  AccumLatDiff();
}

static VecRTwo ComputeTrueLipSiteCoords( VecRTwo rSitei, VecRTwo rSitej, 
					 int iSiteIndex ) // site "i" index in the main site[] array
// ComputeTrueLipSiteCoords computes the true (unwrapped) position vector of sitei 
//given the position of sitej (which precedes sitei in the intralipid architecture)
{
  VecRTwo d;  // inter-site distance vectors, used to work out the real (unwrapped) site coords 
  VCopyTwo( rSitei, siteCoords[ iSiteIndex ] ); 
  VSubTwo( d, rSitei, rSitej ); // compute intersite distance
  VWrapXY( d ); // compute minimum image
  VAddTwo( rSitei, rSitej, d ); // compute unwrapped coordinates
  return rSitei;
}

void ComputeLipLatCom() // works out total "lateral" mass centre coordinates (i.e. x-y coords) - for lateral diffusion measurememnts
{
  int n, sCount, nBeadsLipid, lipSiteCount; // site counter
  int lipCount = 0;
  VecRTwo rSite[ MAX_SITES_PER_LIPID ]; // line above modified to allow compilation with pgi on iridis 
  
  real sFac;

  sFac =  ( 1. / dopcStruct.nBeads ); // ( reasonable ) assumption: every (lipid) site has got the same mass
  /*   printf("*** Warning from ComputeLipLatCom() - only valid for DOPC/DOPE mixtures\n"); */

  DO_SITE { // start scanning all sites - the index "n" will actually be modified inside the lipid switches - see below 
    switch ( site[ n ].lipidSpecies ) {
    case NON_LIPID: break;
    case DOPC: nBeadsLipid = dopcStruct.nBeads; break;
    case DOPE: nBeadsLipid = dopeStruct.nBeads; break;
    default: printf( "Error in lip-com switch: switch ( site[ n ].lipidSpecies ) default reached! Exiting...\n" ); exit(0);
    }
    switch ( site[ n ].lipidSpecies ) {
    case NON_LIPID: break;
    case DOPC: 
    case DOPE: 
      VZeroTwo( xyLipidCom[ lipCount ] ); // initialise com for the current lipid
      for ( sCount = 0; sCount < nBeadsLipid; sCount++ ) {
	lipSiteCount = n + sCount; // 'n' corresponds to a choline_or_amine site; from here, lipSiteCount scans through the rest of the lipid
	switch ( site[ lipSiteCount ].lipidUnit ) {
	case CHOLINE_OR_AMINE: VCopyTwo( rSite[ CHOLINE_OR_AMINE ], siteCoords[ n ] ); break;
	case PHOSPHATE:  rSite[ PHOSPHATE ] = ComputeTrueLipSiteCoords( rSite[ PHOSPHATE ], rSite[ CHOLINE_OR_AMINE ],   lipSiteCount ); break;
	case GLYCEROL:   rSite[ GLYCEROL ]  = ComputeTrueLipSiteCoords( rSite[ GLYCEROL ],  rSite[ PHOSPHATE ], lipSiteCount ); break;
	case ESTER_A:    rSite[ ESTER_A ]   = ComputeTrueLipSiteCoords( rSite[ ESTER_A ],   rSite[ GLYCEROL ],  lipSiteCount ); break;
	case TAIL_A1:    rSite[ TAIL_A1 ] = ComputeTrueLipSiteCoords( rSite[ TAIL_A1 ], rSite[ ESTER_A ], lipSiteCount ); break;
	case TAIL_A2:    rSite[ TAIL_A2 ] = ComputeTrueLipSiteCoords( rSite[ TAIL_A2 ], rSite[ TAIL_A1 ], lipSiteCount ); break;
	case TAIL_A3:    rSite[ TAIL_A3 ] = ComputeTrueLipSiteCoords( rSite[ TAIL_A3 ], rSite[ TAIL_A2 ], lipSiteCount ); break;
	case TAIL_A4:    rSite[ TAIL_A4 ] = ComputeTrueLipSiteCoords( rSite[ TAIL_A4 ], rSite[ TAIL_A3 ], lipSiteCount ); break; 
	case TAIL_A5:    rSite[ TAIL_A5 ] = ComputeTrueLipSiteCoords( rSite[ TAIL_A5 ], rSite[ TAIL_A4 ], lipSiteCount ); break; 
	case ESTER_B:    rSite[ ESTER_B ]   = ComputeTrueLipSiteCoords( rSite[ ESTER_B ],   rSite[ GLYCEROL ],  lipSiteCount ); break;
	case TAIL_B1:    rSite[ TAIL_B1 ] = ComputeTrueLipSiteCoords( rSite[ TAIL_B1 ], rSite[ ESTER_B ], lipSiteCount ); break;
	case TAIL_B2:    rSite[ TAIL_B2 ] = ComputeTrueLipSiteCoords( rSite[ TAIL_B2 ], rSite[ TAIL_B1 ], lipSiteCount ); break;
	case TAIL_B3:    rSite[ TAIL_B3 ] = ComputeTrueLipSiteCoords( rSite[ TAIL_B3 ], rSite[ TAIL_B2 ], lipSiteCount ); break;
	case TAIL_B4:    rSite[ TAIL_B4 ] = ComputeTrueLipSiteCoords( rSite[ TAIL_B4 ], rSite[ TAIL_B3 ], lipSiteCount ); break; 
	case TAIL_B5:    rSite[ TAIL_B5 ] = ComputeTrueLipSiteCoords( rSite[ TAIL_B5 ], rSite[ TAIL_B4 ], lipSiteCount ); break; 
	default: printf( "Error in lip-com switch: DEFAULT reached! Exiting...\n" ); exit(0);
	}
      } // end scanning through current lipid
      // at this point, the array rSite[] contains the unwrapped coords of the current lipid: now its peanuts to get the com!
      for ( lipSiteCount = 1; lipSiteCount <= nBeadsLipid; lipSiteCount++ ) { // scan through individual sites of current lipid
	VVAddTwo( xyLipidCom[ lipCount ],  rSite[ lipSiteCount ] ); // accumulate
      }  // end scanning through current lipid
      VScaleTwo( xyLipidCom[ lipCount ], sFac );
      lipCount++;
      n += nBeadsLipid - 1;
      break;
    default: printf("ERROR in ComputeLipLatCom() switch! Exiting...\n"); exit(0);
    }
  } 
  if ( lipCount != nLipids ) { printf("ERROR in ComputeLipLatCom(): lipCount != nLipids. Exiting...\n"); exit(0); }
}

void UnwrapLipLatCom() 
{
  int n;
  VecRTwo dxy;
  // get the "true" position, i.e. remove PBC wrapping effect

  DO_LIPID {
    VSubTwo( dxy, xyLipidComTrue[ n ], xyLipidCom[ n ] );
    VDivTwo( dxy, dxy, region );
    dxy.x = Nint( dxy.x );
    dxy.y = Nint( dxy.y );
    VMulTwo( dxy, dxy, region );
    VAddTwo( xyLipidComTrue[ n ], xyLipidCom[ n ], dxy );
  }
}

void ComputeAndUnwrapSoluteLatCom() 
{
  VecRTwo dxy;
  VSetTwo( xySoluteCom, siteCoords[ nSites-1 ].x, siteCoords[ nSites-1 ].y ); // usual assumption: solute is the last element of the "site" array 
  VSubTwo( dxy, xySoluteComTrue, xySoluteCom ); 
  VDivTwo( dxy, dxy, region );
  dxy.x = Nint( dxy.x );
  dxy.y = Nint( dxy.y );
  VMulTwo( dxy, dxy, region );
  VAddTwo( xySoluteComTrue, xySoluteCom, dxy );
}

